"""MCP工具实现 - 定义MCP协议的工具"""

import asyncio
from typing import Dict, Any, List, Optional, Union
from datetime import datetime
from loguru import logger

import sys
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

from src.core.context_orchestrator import ContextOrchestrator


class MCPTools:
    """MCP工具集合 - 实现MCP协议的各种工具"""
    
    def __init__(self, orchestrator: ContextOrchestrator):
        self.orchestrator = orchestrator
        self.tools = {}
        self._register_default_tools()
        
        logger.info("MCP工具集合已创建")
    
    def _register_default_tools(self):
        """注册默认工具"""
        # 记忆相关工具
        self.register_tool(
            "memory_search",
            "搜索记忆",
            "在用户的记忆中搜索相关信息",
            {
                "query": {"type": "string", "description": "搜索查询"},
                "user_id": {"type": "string", "description": "用户ID"},
                "memory_type": {"type": "string", "description": "记忆类型", "enum": ["short_term", "long_term", "associative", "semantic", "all"]},
                "limit": {"type": "integer", "description": "结果数量限制", "default": 5}
            },
            ["query", "user_id"],
            self._memory_search
        )
        
        self.register_tool(
            "create_memory",
            "创建记忆",
            "为用户创建新的记忆",
            {
                "content": {"type": "string", "description": "记忆内容"},
                "user_id": {"type": "string", "description": "用户ID"},
                "memory_type": {"type": "string", "description": "记忆类型", "enum": ["episodic", "semantic", "procedural"]},
                "importance": {"type": "number", "description": "重要性评分 (0-1)", "default": 0.5},
                "tags": {"type": "array", "items": {"type": "string"}, "description": "标签列表"}
            },
            ["content", "user_id"],
            self._create_memory
        )
        
        # 知识相关工具
        self.register_tool(
            "knowledge_upload",
            "上传知识文档",
            "上传并处理知识文档",
            {
                "content": {"type": "string", "description": "文档内容"},
                "title": {"type": "string", "description": "文档标题"},
                "user_id": {"type": "string", "description": "用户ID"},
                "metadata": {"type": "object", "description": "文档元数据"}
            },
            ["content", "title"],
            self._knowledge_upload
        )
        
        self.register_tool(
            "knowledge_search",
            "搜索知识库",
            "在知识库中搜索相关文档",
            {
                "query": {"type": "string", "description": "搜索查询"},
                "user_id": {"type": "string", "description": "用户ID"},
                "limit": {"type": "integer", "description": "结果数量限制", "default": 5}
            },
            ["query"],
            self._knowledge_search
        )
        
        # 对话相关工具
        self.register_tool(
            "conversation_create",
            "创建对话",
            "创建新的对话会话",
            {
                "user_id": {"type": "string", "description": "用户ID"},
                "title": {"type": "string", "description": "对话标题"},
                "metadata": {"type": "object", "description": "对话元数据"}
            },
            ["user_id"],
            self._conversation_create
        )
        
        self.register_tool(
            "conversation_search",
            "搜索对话",
            "搜索历史对话",
            {
                "query": {"type": "string", "description": "搜索查询"},
                "user_id": {"type": "string", "description": "用户ID"},
                "limit": {"type": "integer", "description": "结果数量限制", "default": 5}
            },
            ["query"],
            self._conversation_search
        )
        
        # 图谱相关工具
        self.register_tool(
            "graph_query",
            "图谱查询",
            "查询知识图谱",
            {
                "query": {"type": "string", "description": "Cypher查询语句"},
                "parameters": {"type": "object", "description": "Cypher查询参数"},
                "user_id": {"type": "string", "description": "用户ID"},
                "read_only": {"type": "boolean", "description": "是否只读查询", "default": True}
            },
            ["query"],
            self._graph_query
        )
        
        self.register_tool(
            "find_related_concepts",
            "查找相关概念",
            "查找与给定概念相关的其他概念",
            {
                "concept_id": {"type": "string", "description": "概念ID"},
                "concept": {"type": "string", "description": "概念名称（兼容字段）"},
                "depth": {"type": "integer", "description": "搜索深度", "default": 2},
                "limit": {"type": "integer", "description": "结果数量限制", "default": 10}
            },
            [],  # 不设置必需参数，在实现中检查
            self._find_related_concepts
        )
        
        # 用户相关工具
        self.register_tool(
            "get_user_context",
            "获取用户上下文",
            "获取用户的完整上下文信息",
            {
                "user_id": {"type": "string", "description": "用户ID"}
            },
            ["user_id"],
            self._get_user_context
        )
        
        # 系统相关工具
        self.register_tool(
            "system_stats",
            "系统统计",
            "获取系统统计信息",
            {},
            [],
            self._system_stats
        )
        
        self.register_tool(
            "health_check",
            "健康检查",
            "检查系统健康状态",
            {},
            [],
            self._health_check
        )
    
    def register_tool(self, name: str, display_name: str, description: str,
                     input_schema: Dict[str, Any], required: List[str],
                     handler: callable):
        """注册工具"""
        self.tools[name] = {
            "name": name,
            "display_name": display_name,
            "description": description,
            "inputSchema": {
                "type": "object",
                "properties": input_schema,
                "required": required
            },
            "handler": handler
        }
        
        logger.debug(f"注册工具: {name}")
    
    def get_tool_list(self) -> List[Dict[str, Any]]:
        """获取工具列表"""
        return [
            {
                "name": tool["name"],
                "description": tool["description"],
                "inputSchema": tool["inputSchema"]
            }
            for tool in self.tools.values()
        ]
    
    async def call_tool(self, name: str, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """调用工具"""
        try:
            if name not in self.tools:
                return {
                    "isError": True,
                    "content": [{
                        "type": "text",
                        "text": f"工具不存在: {name}"
                    }]
                }
            
            tool = self.tools[name]
            handler = tool["handler"]
            
            # 验证参数
            validation_error = self._validate_arguments(tool["inputSchema"], arguments)
            if validation_error:
                return {
                    "isError": True,
                    "content": [{
                        "type": "text",
                        "text": f"参数验证失败: {validation_error}"
                    }]
                }
            
            # 调用处理器
            result = await handler(arguments)
            
            return {
                "content": [{
                    "type": "text",
                    "text": str(result)
                }]
            }
            
        except Exception as e:
            logger.error(f"调用工具失败: {name} - {e}")
            return {
                "isError": True,
                "content": [{
                    "type": "text",
                    "text": f"工具执行失败: {str(e)}"
                }]
            }
    
    def _validate_arguments(self, schema: Dict[str, Any], arguments: Dict[str, Any]) -> Optional[str]:
        """验证参数"""
        try:
            required_fields = schema.get("required", [])
            properties = schema.get("properties", {})
            
            # 检查必需字段
            for field in required_fields:
                if field not in arguments:
                    return f"缺少必需参数: {field}"
            
            # 检查字段类型（简化版本）
            for field, value in arguments.items():
                if field in properties:
                    field_type = properties[field].get("type")
                    if field_type == "string" and not isinstance(value, str):
                        return f"参数 {field} 应为字符串类型"
                    elif field_type == "integer" and not isinstance(value, int):
                        return f"参数 {field} 应为整数类型"
                    elif field_type == "number" and not isinstance(value, (int, float)):
                        return f"参数 {field} 应为数字类型"
                    elif field_type == "boolean" and not isinstance(value, bool):
                        return f"参数 {field} 应为布尔类型"
                    elif field_type == "array" and not isinstance(value, list):
                        return f"参数 {field} 应为数组类型"
                    elif field_type == "object" and not isinstance(value, dict):
                        return f"参数 {field} 应为对象类型"
            
            return None
            
        except Exception as e:
            return f"参数验证错误: {str(e)}"
    
    # 工具处理器实现
    
    async def _memory_search(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """搜索记忆"""
        try:
            query = arguments["query"]
            user_id = arguments["user_id"]
            memory_type = arguments.get("memory_type")
            limit = arguments.get("limit", 5)

            # 修正参数顺序：MemoryManager.search_memories(user_id, query, memory_types, limit)
            memory_types = [memory_type] if memory_type and memory_type != "all" else None
            results = await self.orchestrator.memory_manager.search_memories(
                user_id,
                query,
                memory_types,
                limit
            )

            return {
                "success": True,
                "results": results,
                "count": len(results) if isinstance(results, list) else 0
            }

        except Exception as e:
            logger.error(f"搜索记忆失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _create_memory(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """创建记忆"""
        try:
            content = arguments["content"]
            user_id = arguments["user_id"]
            memory_type = arguments.get("memory_type", "episodic")
            importance = arguments.get("importance", 0.5)
            tags = arguments.get("tags", [])
            
            memory_id = await self.orchestrator.memory_manager.store_memory(
                user_id, content, memory_type, {"tags": tags}, importance
            )
            
            return {
                "success": True,
                "memory_id": memory_id,
                "message": "记忆创建成功"
            }
            
        except Exception as e:
            logger.error(f"创建记忆失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _knowledge_upload(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """上传知识文档"""
        try:
            content = arguments["content"]
            title = arguments["title"]  # 现在为必填参数
            user_id = arguments.get("user_id")
            metadata = arguments.get("metadata", {})

            document_id = await self.orchestrator.knowledge_processor.upload_document(
                content, title, user_id, metadata
            )

            return {
                "success": True,
                "document_id": document_id,
                "message": "文档上传成功"
            }

        except Exception as e:
            logger.error(f"上传知识文档失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _knowledge_search(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """搜索知识库"""
        try:
            query = arguments["query"]
            user_id = arguments.get("user_id")
            limit = arguments.get("limit", 5)
            
            results = await self.orchestrator.knowledge_processor.search_knowledge(
                query, user_id, limit
            )
            
            return {
                "success": True,
                "results": results,
                "count": len(results)
            }
            
        except Exception as e:
            logger.error(f"搜索知识库失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _conversation_create(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """创建对话"""
        try:
            user_id = arguments["user_id"]
            title = arguments.get("title")
            metadata = arguments.get("metadata", {})
            
            conversation_id = await self.orchestrator.conversation_manager.create_conversation(
                user_id, title, metadata
            )
            
            # 自动添加包含title和metadata的初始消息，使其可被conversation_search搜索到
            metadata_content = ""
            if title:
                metadata_content += f"这是关于'{title}'的对话。"
            
            if metadata:
                # 将metadata转换为可搜索的文本
                metadata_parts = []
                for key, value in metadata.items():
                    if isinstance(value, (str, int, float, bool)):
                        metadata_parts.append(f"{key}: {value}")
                    elif isinstance(value, (list, dict)):
                        metadata_parts.append(f"{key}: {str(value)}")
                
                if metadata_parts:
                    metadata_content += f" 相关信息: {', '.join(metadata_parts)}"
            
            # 如果没有title和metadata，添加默认内容
            if not metadata_content:
                metadata_content = "这是一个新的对话会话。"
            
            # 添加初始消息到对话中，使其可被向量搜索
            await self.orchestrator.conversation_manager.add_message(
                conversation_id=conversation_id,
                role="assistant",
                content=metadata_content,
                metadata={
                    "type": "conversation_metadata",
                    "auto_generated": True,
                    "timestamp": datetime.now().isoformat()
                }
            )
            
            return {
                "success": True,
                "conversation_id": conversation_id,
                "message": "对话创建成功，已添加可搜索的元数据"
            }
            
        except Exception as e:
            logger.error(f"创建对话失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _conversation_search(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """搜索对话"""
        try:
            query = arguments.get("query")
            user_id = arguments.get("user_id")
            limit = arguments.get("limit", 5)
            
            if query:
                results = await self.orchestrator.conversation_manager.search_conversations(
                    query, user_id, limit
                )
            else:
                # 未提供 query 时回退为按用户获取会话列表
                if not user_id:
                    return {"success": False, "error": "缺少 query 或 user_id 参数"}
                conversations = await self.orchestrator.conversation_manager.get_user_conversations(user_id)
                # 简单截断到 limit
                results = conversations[:limit] if isinstance(conversations, list) else conversations
            
            return {
                "success": True,
                "results": results,
                "count": len(results) if isinstance(results, list) else 1
            }
            
        except Exception as e:
            logger.error(f"搜索对话失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _graph_query(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """执行图数据库查询"""
        try:
            query = arguments["query"]
            parameters = arguments.get("parameters", {})

            neo4j = self.orchestrator.db_manager.neo4j_client
            exec_fn = getattr(neo4j, "execute_read_query", None) or getattr(neo4j, "execute_query", None)
            if not exec_fn:
                raise AttributeError("Neo4jClient缺少execute_read_query/execute_query 方法")

            result = await exec_fn(query, parameters)
            return {"success": True, "result": result}
        except Exception as e:
            logger.error(f"图数据库查询失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _find_related_concepts(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """查找与给定概念相关的其他概念（基于图数据库关系检索）"""
        try:
            concept = arguments.get("concept") or arguments.get("concept_id")
            if not concept:
                return {"success": False, "error": "缺少必需参数: concept 或 concept_id"}
            depth = arguments.get("depth", 2)
            limit = arguments.get("limit", 10)

            neo4j = self.orchestrator.db_manager.neo4j_client
            
            # 如果传入的是概念名称而不是concept_id，先查找匹配的概念
            if not concept.startswith('concept_'):
                # 通过名称查找概念
                concepts = await neo4j.find_concepts_by_name(concept, 5)
                if not concepts:
                    return {"success": True, "results": [], "count": 0, "message": f"未找到名称包含'{concept}'的概念"}
                
                # 对找到的每个概念查找相关概念
                all_related = []
                for concept_node in concepts:
                    concept_id = concept_node.get('concept_id')
                    if concept_id:
                        related = await neo4j.find_related_concepts(concept_id, depth, limit)
                        all_related.extend(related)
                
                # 去重并限制结果数量
                seen_ids = set()
                unique_results = []
                for item in all_related:
                    item_id = item.get('concept_id')
                    if item_id and item_id not in seen_ids:
                        seen_ids.add(item_id)
                        unique_results.append(item)
                        if len(unique_results) >= limit:
                            break
                
                return {"success": True, "results": unique_results, "count": len(unique_results)}
            else:
                # 直接使用concept_id查找
                results = await neo4j.find_related_concepts(concept, depth, limit)
                return {"success": True, "results": results, "count": len(results) if isinstance(results, list) else 0}
                
        except Exception as e:
            logger.error(f"查找相关概念失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _get_user_context(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """获取用户上下文"""
        try:
            user_id = arguments["user_id"]
            context = await self.orchestrator.get_user_context(user_id)
            return {"success": True, "context": context}
        except Exception as e:
            # 防止底层MongoDB客户端异常泄漏到工具返回
            logger.warning(f"获取用户上下文时出现问题，将返回空上下文: {e}")
            return {"success": True, "context": {}, "warning": "user context temporarily unavailable"}
    
    async def _system_stats(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """系统统计"""
        try:
            # 兼容 get_stats 与 get_system_stats 两种命名
            get_stats_fn = getattr(self.orchestrator, "get_stats", None) or getattr(self.orchestrator, "get_system_stats", None)
            if not get_stats_fn:
                raise AttributeError("ContextOrchestrator缺少get_stats/get_system_stats 方法")
            stats = await get_stats_fn()

            return {
                "success": True,
                "stats": stats
            }

        except Exception as e:
            logger.error(f"获取系统统计失败: {e}")
            return {"success": False, "error": str(e)}
    
    async def _health_check(self, arguments: Dict[str, Any]) -> Dict[str, Any]:
        """健康检查"""
        try:
            health = await self.orchestrator.health_check()
            
            return {
                "success": True,
                "health": health
            }
            
        except Exception as e:
            logger.error(f"健康检查失败: {e}")
            return {"success": False, "error": str(e)}
    
    def get_tool_info(self, name: str) -> Optional[Dict[str, Any]]:
        """获取工具信息"""
        return self.tools.get(name)
    
    def remove_tool(self, name: str) -> bool:
        """移除工具"""
        if name in self.tools:
            del self.tools[name]
            logger.debug(f"移除工具: {name}")
            return True
        return False
    
    def get_tool_names(self) -> List[str]:
        """获取所有工具名称"""
        return list(self.tools.keys())
    
    def get_tool_count(self) -> int:
        """获取工具数量"""
        return len(self.tools)